package tcpShell

import (
	"fmt"
	"gitee.com/kingzyt/common/countPool"
	"gitee.com/kingzyt/common/log"
	"gitee.com/kingzyt/common/util"
	"net"
	"time"
)

const Tcp_dial_time_out = time.Second * 10

func StartTCPListener(name string, listenPort string, countPool *countPool.CountPool, onAcceptFunc func(conn *net.TCPConn), validIPs []string) (quitSign chan byte, ok bool) {
	if onAcceptFunc == nil {
		panic(fmt.Errorf("name:%s, listenPort:%s, onAcceptFunc is nil", name, listenPort))
	}

	laddr, err := net.ResolveTCPAddr("tcp", ":"+listenPort)
	if err != nil {
		log.Error(name, "ResolveTCPAddr fail, err:%s", err)
		ok = false
		return
	}

	listener, err := net.ListenTCP("tcp", laddr)
	if err != nil {
		log.Error(name, "ListenTCP fail, err:%s", err)
		ok = false
		return
	}

	log.Info(name, log.I, "Start Listening: %s", listener.Addr().String())

	quitSign = make(chan byte, 1)

	go func() {
		defer func() {
			log.Info(name, log.I, "listener [::]:%s quit", listenPort)
			listener.Close()
		}()

	End:
		for {
			select {
			case <-quitSign:
				break End
			default:
				// limit tcp link
				if countPool != nil {
					countPool.GetOne()
				}
				conn, err := listener.AcceptTCP()
				if err != nil {
					log.Error(name, "AcceptTCP fail, err:%s", err)
					if countPool != nil {
						countPool.ReturnOne()
					}
					continue
				}

				localAddr, remoteAddr, ok := GetTCPConnAddr(conn)
				if !ok {
					log.Error(name, "AcceptTCP fail, conn fail, cannot get local/remote addr")
					conn.Close()
					if countPool != nil {
						countPool.ReturnOne()
					}
					continue
				}

				if validIPs != nil {
					valid := false
					inIp := util.GetAddrIP(remoteAddr)
					for _, ip := range validIPs {
						if inIp == ip {
							valid = true
							break
						}
					}
					if !valid {
						log.Warn(name, "invalid ip connect in, reject, ip:%s", inIp)
						conn.Close()
						if countPool != nil {
							countPool.ReturnOne()
						}
						continue
					}
				}

				log.Info(name, log.I, "Accept OK: remote: %s, local: %s", remoteAddr, localAddr)

				go onAcceptFunc(conn)
			}
		}
	}()

	ok = true
	return
}

func StartTCPDialer(name string, raddr string, timeout time.Duration) (dialer *net.TCPConn) {
	conn, err := net.DialTimeout("tcp", raddr, timeout)
	if err != nil {
		log.Error(name, "Dial fail, err:%s", err)
		return
	}

	localAddr, remoteAddr, ok := GetTCPConnAddr(conn)
	if !ok {
		log.Error(name, "Dial fail, conn fail, cannot get local/remote addr")
		conn.Close()
		return
	}

	log.Info(name, log.I, "Dial OK: remote: %s, local: %s", remoteAddr, localAddr)

	dialer, _ = conn.(*net.TCPConn)

	return
}

func GetTCPConnAddr(conn net.Conn) (localAddrStr string, remoteAddrStr string, ok bool) {
	localAddr := conn.LocalAddr()
	if localAddr == nil {
		return
	}
	remoteAddr := conn.RemoteAddr()
	if remoteAddr == nil {
		return
	}

	return localAddr.String(), remoteAddr.String(), true
}

func TCPConnWrite(conn net.Conn, timeout time.Time, pkg []byte, onWriteErr func(error)) {
	pkgLen := len(pkg)

	conn.SetWriteDeadline(timeout)
	for {
		n, err := conn.Write(pkg)
		if err != nil {
			onWriteErr(err)
			return
		}

		if n < pkgLen {
			pkg = pkg[n:]
			pkgLen -= n
		} else {
			break
		}
	}
}
